from typing import Dict
from model.voice_v2pro1.local_api import InitTTS


class TTSModelManager:
    def __init__(self, character_configs: list):
        """
        在构造函数中，预加载所有角色的TTS模型。
        :param character_configs: 从 character_prompt.json 加载的角色配置列表。
        """
        self._models: Dict[str, InitTTS] = {}
        print("【TTS模型管理器】开始初始化...")

        for config in character_configs:
            character_id = config.get("id")
            tts_config = config.get("audio_config")

            if not (character_id and tts_config and
                    "ckpt_checkpoint" in tts_config and
                    "pth_checkpoint" in tts_config):
                print(f"警告：角色 {config.get('name')} 的TTS配置不完整，跳过加载。")
                continue

            print(f"  - 正在为角色ID '{character_id}' ({config.get('name')}) 加载TTS模型...")
            try:
                # 为每个角色创建一个独立的TTS实例
                tts_instance = InitTTS(
                    t2s_weights_path=tts_config["ckpt_checkpoint"],
                    vits_weights_path=tts_config["pth_checkpoint"]
                )
                # 将加载好的模型实例存入字典，用 character_id 作为键
                self._models[character_id] = tts_instance
                print(f"  - 角色ID '{character_id}' 模型加载成功！")
            except Exception as e:
                print(f"  - 错误：为角色ID '{character_id}' 加载模型失败: {e}")

        print(f"【TTS模型管理器】初始化完成。共加载了 {len(self._models)} 个模型。")

    def get_tts_model(self, character_id: str) -> InitTTS | None:
        """
        根据 character_id 获取已经加载好的模型实例。
        :param character_id: 角色ID
        :return: InitTTS 实例或 None
        """
        return self._models.get(character_id)